# coding: utf-8
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle


class Node:
    def __init__(self, x, y):
        self.x = x
        self.y = y

    # Return four neighbor nodes: left, right, up, down
    def neighbors(self):
        return [Node(self.x - 1, self.y),
                Node(self.x + 1, self.y),
                Node(self.x, self.y - 1),
                Node(self.x, self.y + 1)]

    def is_in(self, node_set):
        for node in node_set:
            if self.x == node.x and self.y == node.y:
                return True
        return False

    def __repr__(self):
        return f'({self.x}, {self.y})'



class RandomMap:
    def __init__(self, size=50):
        self.size = size
        self.start = Node(0, 0)
        self.goal = Node(self.size - 1, self.size - 1)
        self.obstacles = []
        self.generate()

    def generate(self):
        # Obstacles in the middle
        for n in range(self.size // 4):
            self.obstacles.append(Node(self.size // 2 - n, self.size // 2 + n))
            self.obstacles.append(Node(self.size // 2 + n, self.size // 2 - n))
        # Random obstacles
        for n in range(self.size // 4):
            x = np.random.randint(0, self.size)
            y = np.random.randint(0, self.size)
            self.obstacles.append(Node(x, y))
            if np.random.rand() < 0.5:  # Vertical
                for i in range(self.size // 4):
                    self.obstacles.append(Node(x, y + i))
            else:  # Horizontal
                for i in range(self.size // 4):
                    self.obstacles.append(Node(x + i, y))

    def is_valid(self, node):
        if node.x < 0 or node.y < 0:
            return False
        if node.x > self.size - 1 or node.y > self.size - 1:
            return False
        for obstacle in self.obstacles:
            if node.x == obstacle.x and node.y == obstacle.y:
                return False
        return True


def distance(a, b):
    return abs(a.x - b.x) + abs(a.y - b.y)


def g_func(node, start):
    return distance(node, start)


def h_func(node, goal):
    return distance(node, goal)


def A_star(graph, start, goal):
    frontier = [start]
    close_set = []
    f_score = {start: h_func(start, goal)}
    came_from = {start: None}
    cost_so_far = {start: 0}
    while frontier:
        n = min(frontier, key=lambda node: f_score[node])
        if n.x == goal.x and n.y == goal.y:
            break
        close_set.append(n)
        for m in n.neighbors():
            if graph.is_valid(m):
                if m.is_in(close_set):
                    continue
                else:
                    new_cost = cost_so_far[n] + distance(n, m)
                    if m not in cost_so_far or new_cost < cost_so_far[m]:
                        cost_so_far[m] = new_cost
                        f_score[m] = new_cost + h_func(m, goal)
                        frontier.append(m)
                        came_from[m] = n
        frontier.remove(n)
        draw_graph(graph, frontier, close_set)

    return came_from, cost_so_far


def draw_graph(graph, frontier=None, close_set=None):
    plt.clf()
    plt.plot()
    plt.xlim(0, graph.size)
    plt.ylim(0, graph.size)
    plt.axis('equal')
    plt.axis('off')
    ax = plt.gca()
    for x in range(graph.size):
        for y in range(graph.size):
            if graph.is_valid(Node(x, y)):
                rect = Rectangle((x, y), height=1, width=1, edgecolor='gray', facecolor='w')
            else:
                rect = Rectangle((x, y), height=1, width=1, edgecolor='gray', facecolor='gray')
            ax.add_patch(rect)
    if frontier:
        for node in frontier:
            rect = Rectangle((node.x, node.y), height=1, width=1, edgecolor='gray', facecolor='blue')
            ax.add_patch(rect)
    if close_set:
        for node in close_set:
            rect = Rectangle((node.x, node.y), height=1, width=1, edgecolor='gray', facecolor='orange')
            ax.add_patch(rect)
    rect = Rectangle((0, 0), height=1, width=1, color='green')
    ax.add_patch(rect)
    rect = Rectangle((0, graph.size - 1), height=1, width=1, color='red')
    ax.add_patch(rect)
    plt.pause(0.001)


rand_map = RandomMap(20)
draw_graph(rand_map)
A_star(graph=rand_map, start=Node(0, 0), goal=Node(0, rand_map.size-1))


# plt.show()
